
import numpy as np
import cv2
import time
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from ai.base.base import BaseFunc


class OpenCVOp:
    def __init__(self):
        self.basepath = "J:/av"

    def test_image(self):
        img = np.array([
			[[255,0,0], [0,255,0], [0,0,255]],
			[[255,255,0], [255,0,255], [0,255,255]],
			[[255,255,255], [128,128,128],[0,0,0]],],
			dtype = np.uint8)
        # plt.imsave("img_pyplot.jpg", img)
		# cv2.imwrite("img_cv2.jpg", img)
        cv2.imshow("img", img)
        cv2.waitKey(0)

    def test_imgbase(self, name):
        imgpath = self.basepath + "/" + name
        print("Open image %s" % imgpath)
        color_img = cv2.imread(imgpath)
        print(color_img.shape)
        gray_img = cv2.imread(imgpath, cv2.IMREAD_GRAYSCALE)
        print(gray_img.shape)
        cv2.imwrite(self.basepath + "/test_grayscale.jpg", gray_img)
        reload_grayscale = cv2.imread(self.basepath + "/test_grayscale.jpg")
        print(reload_grayscale.shape)
        cv2.imwrite(self.basepath + "/test_imwrite.jpg", color_img, (cv2.IMWRITE_JPEG_QUALITY, 20))
        cv2.imwrite(self.basepath + "/test_imwrite.png", color_img, (cv2.IMWRITE_PNG_COMPRESSION, 9))

    def test_resize(self, name):
        img = cv2.imread(self.basepath + "/" + name)
        img_200x200 = cv2.resize(img, (200, 200))
        img_200x300 = cv2.resize(img, (0, 0), fx=0.5, fy=0.5, interpolation = cv2.INTER_NEAREST)
        img_300x300 = cv2.copyMakeBorder(img, 50, 50, 0, 0, cv2.BORDER_CONSTANT, value = (255, 0, 0))

        patch_tree = img[800:1000, -900:-700]
        cv2.imshow("img_cropped", patch_tree)
        cv2.imshow("img_200x200", img_200x200)
        cv2.imshow("img_200x300", img_200x300)
        cv2.imshow("img_300x300", img_300x300)
        cv2.waitKey(0)
        # cv2.imwrite(self.basepath + "/cropped_tree.jpg", patch_tree)
        # cv2.imwrite(self.basepath + "/resized+200x200.jpg", img_200x200)
        # cv2.imwrite(self.basepath + "/resized+200x300.jpg", img_200x300)
        # cv2.imwrite(self.basepath + "/bordered_300x300.jpg", img_300x300)

    def test_hsv(self, name):
        img = cv2.imread(self.basepath + "/" + name)
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        turn_green_hsv = img_hsv.copy()
        turn_green_hsv[:, :, 0] = (turn_green_hsv[:, :, 0] + 15) % 180
        turn_green_img = cv2.cvtColor(turn_green_hsv, cv2.COLOR_HSV2BGR)

        colorless_hsv = img_hsv.copy()
        colorless_hsv[:,:,1] = 0.5 * colorless_hsv[:,:,1]
        colorless_img = cv2.cvtColor(colorless_hsv, cv2.COLOR_HSV2BGR)

        darker_hsv = img_hsv.copy()
        darker_hsv[:,:,2] = 0.5 * darker_hsv[:,:,2]
        darker_img = cv2.cvtColor(darker_hsv, cv2.COLOR_HSV2BGR)

        cv2.imshow("img", img)
        cv2.imshow("turn green", turn_green_img)
        cv2.imshow("colorless", colorless_img)
        cv2.imshow("darker", darker_img)
        cv2.waitKey(0)

    def gamma_trans(self, img, gamma):
        gamma_table = [np.power(x/255.0, gamma) * 255.0 for x in range(256)]
        gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)
        return cv2.LUT(img, gamma_table)

    def test_gamma(self, name):
        path = self.basepath + "/" + name
        img = cv2.imread(path)
        hist_b = cv2.calcHist([img], [0], None, [256], [0, 256])
        hist_g = cv2.calcHist([img], [1], None, [256], [0,256])
        hist_r = cv2.calcHist([img], [2], None, [256], [0, 256])
        img_corrected = self.gamma_trans(img, 0.2)

        cv2.imshow("img", img)
        cv2.imshow("corrected", img_corrected)
        cv2.waitKey(0)

        hist_b_corrected = cv2.calcHist([img_corrected], [0], None, [256], [0, 256])
        hist_g_corrected = cv2.calcHist([img_corrected], [1], None, [256], [0, 256])
        hist_r_corrected = cv2.calcHist([img_corrected], [2], None, [256], [0, 256])

        fig = plt.figure()
        pix_hists = [
            [hist_b, hist_g, hist_r],
            [hist_b_corrected, hist_g_corrected, hist_r_corrected]
        ]
        pix_vals = range(256)
        for sub_plt, pix_hist in zip([121, 122], pix_hists):
            ax = fig.add_subplot(sub_plt, projection='3d')
            for c, z, channel_hist in zip(['b', 'g', 'r'], [20, 10, 0], pix_hist):
                cs = [c] * 256
                ax.bar(pix_vals, channel_hist, zs = z, zdir='y', color = cs, 
                    alpha = 0.618, edgecolor = 'none', lw = 0)
            ax.set_xlabel('Pixel Values')
            ax.set_xlim([0, 256])
            ax.set_ylabel('Counts')
            ax.set_zlabel('Channels')
        plt.show()

    def test_affine(self, name):
        path = self.basepath + "/" + name
        img = cv2.imread(path)
        M_crop_elephant = np.array([
            [1.6, 0, -150],
            [0, 1.6, -240]
            ], dtype = np.float32)
        img_elephant = cv2.warpAffine(img, M_crop_elephant, (333, 640))

        theta = 15 * np.pi / 180
        M_shear = np.array([
                [1, np.tan(theta), 0],
                [0, 1, 0]
            ], dtype = np.float32)
        img_sheared = cv2.warpAffine(img, M_shear, (333, 640))

        M_rotate = np.array([
                [np.cos(theta), -np.sin(theta), 0],
                [np.sin(theta), np.cos(theta), 0]
            ], dtype= np.float32)
        img_rotated = cv2.warpAffine(img, M_rotate, (333, 640))

        M = np.array([
                [1, 1.5, -400],
                [0.5, 2, -100]
            ], dtype = np.float32)
        img_transformed = cv2.warpAffine(img, M, (333, 640))

        cv2.imshow("img", img)
        cv2.imshow("crop", img_elephant)
        cv2.imshow("sheared", img_sheared)
        cv2.imshow("rotated", img_rotated)
        cv2.imshow("transform", img_transformed)
        cv2.waitKey(0)

class OpenCVVideo:
    def __init__(self):
        self.interval = 0
        self.num_frames = 500
        self.out_fps = 24
        self.basepath = "J:/av/"

    def cap_save(self, save):
        cap = cv2.VideoCapture(0)
        size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))

        print("video capture size ", size)
        video = cv2.VideoWriter(
                self.basepath + save + ".avi",
                cv2.VideoWriter_fourcc('M', 'P', '4', '2'),
                    self.out_fps,
                    size
            )

        for i in range(42):
            cap.read()

        try:
            for i in range(self.num_frames):
                _, frame = cap.read()
                video.write(frame)
                # filename = '{:0>6d}.png'.format(i)
                # cv2.imwrite(fliename, frame)
                print('Frame {} is captured.'.format(i))
                time.sleep(self.interval)
        except KeyboardInterrupt:
            print('Stopped. {}/{} frames captured'.format(i, self.num_frames))

        video.release()
        cap.release()

def test_gaussian_blur():
    b = BaseFunc()
    image = cv2.imread("../imgs/1.jpg");
    outimage = b.gaussian_blur(image, 15)
    b.show_window('1', image)
    b.show_window('2', outimage)
    cv2.waitKey(0)

def test_equalize():
    b = BaseFunc()
    image = cv2.imread("../imgs/1.jpg")
    outimage = b.equalize(image)
    b.show_window('1', outimage)
    cv2.waitKey(0)

if __name__ == "__main__":
    ocv = OpenCVOp()
	# ocv.test_image()
    # ocv.test_imgbase("1.JPG")
    # ocv.test_resize("1.JPG")
    # ocv.test_hsv("1.JPG")
    # ocv.test_gamma("1.jpg")
    # ocv.test_affine("1.jpg")

    # ocvi = OpenCVVideo()
    # ocvi.cap_save("hello")

    # test_gaussian_blur()
    test_equalize()

